import functools

import pandas as pd
import numpy as np

def group_test():
    df = pd.DataFrame({'A' : ['foo', 'bar', 'foo', 'bar',
                              'foo', 'bar', 'foo', 'foo'],
                       'B' : ['one', 'one', 'two', 'three',
                              'two', 'two', 'one', 'three'],
                       'C' : np.random.randn(8),
                       'D' : np.random.randn(8)})

    print("df=\n", df)
    df_groupA = df.groupby('A').sum()
    print("df_groupby A = \n", df_groupA)
    print("df.groupby['A', 'B']", df.groupby(['A','B']).sum())


def group_apply_test():
    df = pd.DataFrame({'animal': 'cat dog cat fish dog cat cat'.split(),
                       'size': list('SSMMMLL'),
                       'weight': [8, 10, 11, 1, 20, 12, 12],
                       'adult': [False] * 5 + [True] * 2})

    print("df:\n", df)
    print("idxmax:\n", df['size'][df['weight'].idxmax()])
    print("result:\n",
        df.groupby('animal').apply(lambda subf: subf['size'][subf['weight'].idxmax()])
    )

def group_get_group_test():
    df = pd.DataFrame({'animal': 'cat dog cat fish dog cat cat'.split(),
                       'size': list('SSMMMLL'),
                       'weight': [8, 10, 11, 1, 20, 12, 12],
                       'adult': [False] * 5 + [True] * 2})
    gb = df.groupby(['animal'])
    print("get cat:\n", gb.get_group('cat'))

    def GrowUp(x):
        avg_weight = sum(x[x['size'] == 'S'].weight * 1.5)
        avg_weight += sum(x[x['size'] == 'M'].weight * 1.25)
        avg_weight += sum(x[x['size'] == 'L'].weight)
        avg_weight /= len(x)
        return pd.Series(['L', avg_weight, True], index=['size', 'weight', 'adult'])

    expected_df = gb.apply(GrowUp)
    print("expected_df:\n", expected_df)

def expanding_apply_test():
    s = pd.Series(i/100.0 for i in range(1,11))
    print("s:\n", s)
    def CumRet(x,y):
        return x * (1 + y)

    def Red(x):
        return functools.reduce(CumRet,x,1.0)
    se = s.expanding().apply(Red)
    print("expanding s:\n", se)

def replace_test():
    df = pd.DataFrame({'A': [1, 1, 2, 2], 'B': [1, -1, 1, 2]})
    print(df)
    gb = df.groupby('A')
    def replace(g):
        mask = g < 0
        print(g)
        g.loc[mask] = g[~mask].mean()
        print(g)
        return g

    print(gb.transform(replace))

def aggregate_test():
    df = pd.DataFrame({'code': ['foo', 'bar', 'baz'] * 2,
                       'data': [0.16, -0.21, 0.33, 0.45, -0.59, 0.62],
                       'flag': [False, True] * 3})
    print("df:\n", df)
    code_groups = df.groupby('code')
    print("code_groups:\n", code_groups[['data']])
    agg_n_sort_order = code_groups[['data']].transform(sum).sort_values(by='data')
    sorted_df = df.ix[agg_n_sort_order.index]
    print("sorted_df:\n", sorted_df)


def aggregate2_test():
    rng = pd.date_range(start="2014-10-07", periods=10, freq='2min')
    ts = pd.Series(data=list(range(10)), index=rng)
    print("ts:\n", ts)
    def MyCust(x):
        if len(x) > 2:
           return x[1] * 1.234
        return pd.NaT

    mhc = {'Mean': np.mean, 'Max': np.max, 'Custom': MyCust}
    print("result:\n", ts.resample("5min").apply(mhc))


def group_count_assign():
    df = pd.DataFrame({'Color': 'Red Red Red Blue'.split(),
                       'Value': [100, 150, 50, 50]})

    # s = df.groupby(['Color']).transform(len)
    # print(s)
    df['Counts'] = df.groupby(['Color']).transform(len)
    print(df)

def shift_test():
    df = pd.DataFrame(
        {u'line_race': [10, 10, 8, 10, 10, 8],
         u'beyer': [99, 102, 103, 103, 88, 100]},
        index=[u'Last Gunfighter', u'Last Gunfighter', u'Last Gunfighter',
               u'Paynter', u'Paynter', u'Paynter'])
    print("df:\n", df)
    df['beyer_shifted'] = df.groupby(level=0)['beyer'].shift(1)
    print("df:\n", df)

def maximum_test():
    df = pd.DataFrame({'host': ['other', 'other', 'that', 'this', 'this'],
                       'service': ['mail', 'web', 'mail', 'mail', 'web'],
                       'no': [1, 2, 1, 2, 1]}).set_index(['host', 'service'])
    print("df:\n", df)
    mask = df.groupby(level=0).agg('idxmax')
    print("mask:\n", mask)
    df_count = df.loc[mask['no']].reset_index()
    print("df_count:\n", df_count)

def grouping_test():
    df = pd.DataFrame([0, 1, 0, 1, 1, 1, 0, 1, 1], columns=['A'])
    print("A.shift:\n", df.A.shift())
    print(df.A.groupby((df.A != df.A.shift()).cumsum()).groups)
    print(df.A.groupby((df.A != df.A.shift()).cumsum()).cumsum())

if __name__ == '__main__':
    # group_test()
    # group_apply_test()
    # group_get_group_test()
    # expanding_apply_test()
    # replace_test()
    # aggregate_test()
    # aggregate2_test()
    # group_count_assign()
    # shift_test()
    # maximum_test()
    grouping_test()